#include <err.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>
#include <getopt.h>
#include <limits.h>
#include <stdint.h>
#include <inttypes.h>

#include "../util.h"
#include "../getopt_helper.h"

#define _GNU_SOURCE
#include <sched.h>

#define CPU_DIR "/sys/devices/system/cpu"

enum {
	ALL = GETOPT_LONG_ENUM_INIT,
	IGNORE,
};

static struct option const LONG_OPTIONS[] = {
	{ "all", no_argument, NULL, ALL },
	{ "ignore", required_argument, NULL, IGNORE },
	GETOPT_END_DECL
};

static const char usage_str[] =
    "Usage: %s [OPTION]...\n"
    "Print the number of processing units available to the current process,\n"
    "which may be less than the number of online processors\n"
    "\n"
    "      --all      print the number of installed processors\n"
    "      --ignore=N if possible, exclude N processing units\n"
    "      --help     display this help and exit\n"
    "      --version  output version information and exit";

typedef uint32_t nproc_t;

static nproc_t sys_nproc(void);
static nproc_t sched_nproc(void);

int
main(int argc, char *argv[])
{
	int ch;

	nproc_t nproc, ignore = 0;

	nproc_t (*nproc_func)(void) = sched_nproc;

	while ((ch = getopt_long(argc, argv, "", LONG_OPTIONS, NULL)) != -1) {
		switch (ch) {
		case ALL:
			nproc_func = sys_nproc;
			break;

		case IGNORE:
			STR2UINT(ignore, optarg, UINT_MAX, {
				err(!0, "ignore");
			});
			break;

			GETOPT_END_CASE
		}
	}

	if (*(argv += optind)) errx(!0, "extra operand '%s'", *argv);

	nproc = nproc_func();

	if (ignore < nproc)
		nproc -= ignore;
	else
		nproc = 1;

	printf("%" PRIu32 "\n", nproc);

	return 0;
}

static nproc_t
sys_nproc(void)
{
	DIR  *cpud;
	char *cpuid;

	struct dirent *de;

	nproc_t nproc = 0;

	if (!!(cpud = opendir(CPU_DIR))) {
		while (!!(de = readdir(cpud)))
			if (!!(cpuid = strstr(de->d_name, "cpu"))
			    && isdigit(cpuid[strlen(cpuid) - 1]))
				nproc++;

		closedir(cpud);
	}

	return N_OR_1(nproc);
}

static nproc_t
sched_nproc(void)
{
	uint8_t bit;
	uint8_t set[sizeof(cpu_set_t)];

	size_t	i;
	nproc_t nproc = 0;

	if (!sched_getaffinity(0, sizeof(set), (void *)set))
		/* from musl libc */
		for (i = 0; i < sizeof(cpu_set_t); i++)
			for (bit = 0; bit < 8; bit++)
				if (set[i] & (1 << bit)) nproc++;

	return N_OR_1(nproc);
}
